package com.kapcb.async.wrapper;

import com.kapcb.async.callback.DefaultCallback;
import com.kapcb.async.callback.ICallback;
import com.kapcb.async.callback.IWorker;
import com.kapcb.async.exception.SkippedException;
import com.kapcb.async.worker.ResultState;
import com.kapcb.async.worker.WorkResult;
import lombok.EqualsAndHashCode;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * <a>Title: WorkerWrapper </a>
 * <a>Author: Kapcb <a>
 * <a>Description: 对每个 worker 及 callback 进行包装，一对一 <a>
 *
 * @author Kapcb
 * @version 1.0
 * @date 2024/5/10 20:27
 * @since 1.0
 */
//@Slf4j
@ToString
@EqualsAndHashCode
public class WorkerWrapper<T, R> {

    /**
     * 任务执行完成
     */
    private static final int FINISH = 1;

    /**
     * 任务执行出错
     */
    private static final int ERROR = 2;

    /**
     * 任务执行中
     */
    private static final int WORKING = 3;

    /**
     * 任务初始化
     */
    private static final int INIT = 0;

    /**
     * 该 wrapper 的唯一标识
     */
    private String id;

    /**
     * worker 将要处理的任务的参数
     */
    private T param;

    private IWorker<T, R> worker;

    private ICallback<T, R> callback;

    /**
     * 在自己后面的 wrapper，如果没有，自己就是末尾；如果有一个，就是串行；如果有多个，有几个就需要开几个线程
     * 通过 must 字段来控制是否依赖项必须完成
     * 1
     * --------3
     * 2
     * 1、2 执行完毕后才能执行3
     */
    private List<WorkerWrapper<?, ?>> nextWrappers;

    /**
     * 依赖的 wrappers，有 2 种情况，1:必须依赖的全部完成后，才能执行自己 2:依赖的任何一个、多个完成了，就可以执行自己
     * 通过 must 字段来控制是否依赖项必须完成
     * 1
     * -------3
     * 2
     * 1、2 执行完毕后才能执行 3
     */
    private List<DependencyWrapper> dependencyWrappers;

    /**
     * 标记该事件是否已经被处理过了，譬如已经超时返回 false 了，后续 RPC 又收到返回值了，则不再二次回调
     * 经试验,volatile 并不能保证"同一毫秒"内,多线程对该值的修改和拉取
     * 1-finish, 2-error, 3-working
     */
    private AtomicInteger state = new AtomicInteger(0);

    /**
     * 该 map 存放所有 wrapper 的 id 和 wrapper 的映射
     */
    private Map<String, WorkerWrapper<?, ?>> forParamUseWrappers;

    /**
     * 也是个钩子变量，用来存临时的结果
     */
    private volatile WorkResult<R> workResult = WorkResult.defaultResult();

    /**
     * 是否在执行自己前，去校验 nextWrapper 的执行结果
     * 1   4
     * -------3
     * 2
     * 如这种在 4 执行前，可能 3 已经执行完毕了（被 2 执行完后触发的），那么 4 就没必要执行了。
     * 注意，该属性仅在 nextWrapper 数量 <=1 时有效，>1 时的情况是不存在的
     */
    private volatile boolean needCheckNextWrapperResult = true;

    private WorkerWrapper(String id, IWorker<T, R> worker, T param, ICallback<T, R> callback) {
        if (worker == null) {
            throw new NullPointerException("async worker is null");
        }
        this.worker = worker;
        this.param = param;
        this.id = id;
        // 允许不设置回调
        if (callback == null) {
            callback = new DefaultCallback<>();
        }
        this.callback = callback;
    }

    /**
     * 开始工作
     *
     * @param executorService
     * @param workerWrapper       代表这次 work 是由哪个上游 wrapper 发起的
     * @param remainTime
     * @param forParamUseWrappers
     */
    private void work(ExecutorService executorService, WorkerWrapper<?, ?> workerWrapper, long remainTime, Map<String, WorkerWrapper<?, ?>> forParamUseWrappers) {
        this.forParamUseWrappers = forParamUseWrappers;
        // 将自己放到所有 wrapper 的集合中
        this.forParamUseWrappers.put(id, this);
        long now = System.currentTimeMillis();

        // 总的已经超时了，就快速失败，进行下一个
        if (remainTime <= 0) {
            fastFail(INIT, null);
            beginNext(executorService, now, remainTime);
            return;
        }


        // 如果自己已经执行过了
        // 可能有多个依赖，其中的一个依赖已经执行完了，并且自己也已开始执行或执行完毕。当另一个依赖执行完毕，又进来该方法时，就不重复处理了
        if (getState() == FINISH || getState() == ERROR) {
            beginNext(executorService, now, remainTime);
            return;
        }

        // 如果在执行前需要校验nextWrapper的状态
        if (needCheckNextWrapperResult) {
            // 如果自己的next链上有已经出结果或已经开始执行的任务了，自己就不用继续了
            if (!checkIsNullResult()) {
                fastFail(INIT, new SkippedException());
                beginNext(executorService, now, remainTime);
                return;
            }
        }

        // 如果没有任何依赖，说明自己就是第一批要执行的
        if (dependencyWrappers == null || dependencyWrappers.isEmpty()) {
            fire();
            beginNext(executorService, now, remainTime);
            return;
        }

        // 如果有前方依赖，存在两种情况
        // 一种是前面只有一个wrapper。即 A  ->  B
        // 一种是前面有多个wrapper。A C D ->   B。需要A、C、D都完成了才能轮到B。但是无论是A执行完，还是C执行完，都会去唤醒B。
        // 所以需要B来做判断，必须A、C、D都完成，自己才能执行

        // 只有一个依赖
        if (dependencyWrappers.size() == 1) {
            doDependsOneJob(workerWrapper);
            beginNext(executorService, now, remainTime);
        } else {
            //有多个依赖时
            doDependsJobs(executorService, dependencyWrappers, workerWrapper, now, remainTime);
        }
    }

    public void work(ExecutorService executorService, long remainTime, Map<String, WorkerWrapper<?, ?>> forParamUseWrappers) {
        work(executorService, null, remainTime, forParamUseWrappers);
    }

    /**
     * 总控制台超时，停止所有任务
     */
    public void stopNow() {
        if (getState() == INIT || getState() == WORKING) {
            fastFail(getState(), null);
        }
    }

    /**
     * 判断自己下游链路上，是否存在已经出结果的或已经开始执行的
     *
     * @return 如果没有返回 true，如果有返回 false
     */
    private boolean checkNextWrapperResult() {
        // 如果自己就是最后一个，或者后面有并行的多个，就返回true
        if (nextWrappers == null || nextWrappers.size() != 1) {
            return getState() == INIT;
        }

        WorkerWrapper<?, ?> nextWorker = nextWrappers.get(0);
        boolean state = nextWorker.getState() == INIT;
        // 继续校验自己的后置节点的状态
        return state && nextWorker.checkNextWrapperResult();
    }

    /**
     * 执行自己的任务，具体的执行是在另一个线程里，但判断阻塞超时是在 work 线程
     */
    private void fire() {
        // 阻塞获取结果
        workResult = workerDoJob();
    }

    /**
     * 快速失败
     *
     * @param expect    期望值
     * @param exception 异常
     * @return true：快速失败成功、false：快速失败失败
     */
    private boolean fastFail(int expect, Exception exception) {
        // 试图将它从 expect 状态修改为 error
        if (!compareAndSetState(expect, ERROR)) {
            return false;
        }

        // 尚未处理过结果
        if (checkIsNullResult()) {
            if (exception == null) {
                workResult = defaultResult();
            } else {
                workResult = defaultExResult(exception);
            }
        }

        callback.result(false, param, workResult);
        return true;
    }

    /**
     * 执行下一个任务
     *
     * @param executorService
     * @param now
     * @param remainTime
     */
    private void beginNext(ExecutorService executorService, long now, long remainTime) {
        // 花费的时间
        long costTime = System.currentTimeMillis() - now;
        if (nextWrappers == null) {
            return;
        }
        int taskCnt = nextWrappers.size();
        if (taskCnt == 1) {
            nextWrappers.get(0).work(executorService, WorkerWrapper.this, remainTime - costTime, forParamUseWrappers);
            return;
        }

        CompletableFuture[] futures = new CompletableFuture[nextWrappers.size()];
        for (int i = 0; i < taskCnt; i++) {
            int finalI = i;
            futures[i] = CompletableFuture.runAsync(() -> nextWrappers.get(finalI).work(executorService, WorkerWrapper.this, remainTime - costTime, forParamUseWrappers), executorService);
        }

        try {
            CompletableFuture.allOf(futures).get(remainTime - costTime, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
//            log.error("get task result error, err msg : {}", e.getMessage());
        }
    }

    private void doDependsOneJob(WorkerWrapper<?, ?> dependWrapper) {
        if (ResultState.TIMEOUT == dependWrapper.getWorkResult().getResultState()) {
            workResult = defaultResult();
            fastFail(INIT, null);
        } else if (ResultState.EXCEPTION == dependWrapper.getWorkResult().getResultState()) {
            workResult = defaultExResult(dependWrapper.getWorkResult().getException());
            fastFail(INIT, null);
        } else {
            // 前面任务结束了，该自己执行任务了
            fire();
        }
    }

    private synchronized void doDependsJobs(ExecutorService executorService, List<DependencyWrapper> dependencyWrappers, WorkerWrapper<?, ?> workerWrapper, long now, long remainTime) {
        // 如果当前任务已经完成了，依赖的其他任务拿到锁再进来时，不需要执行下面的逻辑了
        if (getState() != INIT) {
            return;
        }

        boolean nowDependIsMust = false;
        // 创建必须完成的上游 wrapper 集合
        Set<DependencyWrapper> mustWrapper = new HashSet<>();
        for (DependencyWrapper dependencyWrapper : dependencyWrappers) {
            if (dependencyWrapper.isMust()) {
                mustWrapper.add(dependencyWrapper);
            }
            if (dependencyWrapper.getDependWrapper().equals(workerWrapper)) {
                nowDependIsMust = dependencyWrapper.isMust();
            }
        }

        // 如果全部是不必须的条件，那么只要到了这里，就执行自己
        if (mustWrapper.isEmpty()) {
            if (ResultState.TIMEOUT == workerWrapper.getWorkResult().getResultState()) {
                fastFail(INIT, null);
            } else {
                fire();
            }
            beginNext(executorService, now, remainTime);
            return;
        }

        // 如果存在需要必须完成的，且fromWrapper不是必须的，就什么也不干
        if (!nowDependIsMust) {
            return;
        }

        boolean existNoFinish = false;
        boolean hasError = false;

        for (DependencyWrapper dependencyWrapper : mustWrapper) {
            WorkerWrapper<?, ?> workWrapper = dependencyWrapper.getDependWrapper();
            WorkResult<?> tempWorkResult = workWrapper.getWorkResult();
            if (workWrapper.getState() == INIT || workWrapper.getState() == WORKING) {
                existNoFinish = true;
                break;
            }
            if (tempWorkResult.getResultState() == ResultState.TIMEOUT) {
                workResult = defaultResult();
                hasError = true;
                break;
            }
            if (tempWorkResult.getResultState() == ResultState.EXCEPTION) {
                workResult = defaultExResult(workWrapper.getWorkResult().getException());
                hasError = true;
                break;
            }
        }

        // 只要有失败的
        if (hasError) {
            fastFail(INIT, null);
            beginNext(executorService, now, remainTime);
            return;
        }

        // 如果上游都没有失败，分为两种情况，一种是都finish了，一种是有的在working
        // 都finish的话
        if (!existNoFinish) {
            // 上游都finish了，进行自己
            fire();
            beginNext(executorService, now, remainTime);
            return;
        }

    }

    /**
     * 具体的单个 worker 执行任务
     *
     * @return {@link WorkResult}
     */
    private WorkResult<R> workerDoJob() {
        // 避免重复执行
        if (!checkIsNullResult()) {
            return workResult;
        }
        try {
            // 如果已经不是 INIT 状态了，说明正在被执行或已执行完毕。这一步很重要，可以保证任务不被重复执行
            if (!compareAndSetState(INIT, WORKING)) {
                return workResult;
            }

            callback.begin();

            //执行耗时操作
            R resultValue = worker.action(param, forParamUseWrappers);

            // 如果状态不是在 WORKING，说明别的地方已经修改了，直接返回
            if (!compareAndSetState(WORKING, FINISH)) {
                return workResult;
            }

            workResult.setResultState(ResultState.SUCCESS);
            workResult.setResult(resultValue);
            // 成功回调
            callback.result(true, param, workResult);

            return workResult;

        } catch (Exception exception) {
            // 避免重复回调
            if (!checkIsNullResult()) {
                return workResult;
            }
            fastFail(WORKING, exception);
            return workResult;
        }
    }

    private boolean compareAndSetState(int expect, int update) {
        return this.state.compareAndSet(expect, update);
    }

    private boolean checkIsNullResult() {
        return ResultState.DEFAULT == workResult.getResultState();
    }

    private WorkResult<R> defaultResult() {
        workResult.setResultState(ResultState.TIMEOUT);
        workResult.setResult(worker.defaultValue());
        return workResult;
    }

    private WorkResult<R> defaultExResult(Exception ex) {
        workResult.setResultState(ResultState.EXCEPTION);
        workResult.setResult(worker.defaultValue());
        workResult.setException(ex);
        return workResult;
    }

    private int getState() {
        return state.get();
    }

    private void setNeedCheckNextWrapperResult(boolean needCheckNextWrapperResult) {
        this.needCheckNextWrapperResult = needCheckNextWrapperResult;
    }

    public WorkResult<R> getWorkResult() {
        return workResult;
    }

    public List<WorkerWrapper<?, ?>> getNextWrappers() {
        return nextWrappers;
    }

    public void setParam(T param) {
        this.param = param;
    }

    private void addDepend(DependencyWrapper dependencyWrapper) {
        if (dependencyWrappers == null) {
            dependencyWrappers = new ArrayList<>();
        }
        // 如果依赖的是重复的同一个，就不重复添加了
        for (DependencyWrapper wrapper : dependencyWrappers) {
            if (wrapper.equals(dependencyWrapper)) {
                return;
            }
        }
        dependencyWrappers.add(dependencyWrapper);
    }

    private void addDepend(WorkerWrapper<?, ?> workerWrapper, boolean must) {
        addDepend(new DependencyWrapper(workerWrapper, must));
    }

    private void addNext(WorkerWrapper<?, ?> workerWrapper) {
        if (nextWrappers == null) {
            nextWrappers = new ArrayList<>();
        }
        //避免添加重复
        for (WorkerWrapper<?, ?> nextWrapper : nextWrappers) {
            if (nextWrapper.equals(workerWrapper)) {
                return;
            }
        }
        nextWrappers.add(workerWrapper);
    }

    private void addNextWrappers(List<WorkerWrapper<?, ?>> wrappers) {
        if (wrappers == null || wrappers.isEmpty()) {
            return;
        }
        for (WorkerWrapper<?, ?> wrapper : wrappers) {
            addNext(wrapper);
        }
    }

    private void addDependWrappers(List<DependencyWrapper> dependWrappers) {
        if (dependWrappers == null || dependWrappers.isEmpty()) {
            return;
        }
        for (DependencyWrapper dependWrapper : dependWrappers) {
            addDepend(dependWrapper);
        }
    }

    public static class Builder<T, R> {

        /**
         * 该wrapper的唯一标识
         */
        private String id = UUID.randomUUID().toString();

        /**
         * worker 将来要处理的 param
         */
        private T param;

        private IWorker<T, R> worker;

        private ICallback<T, R> callback;

        /**
         * 自己后面的所有
         */
        private List<WorkerWrapper<?, ?>> nextWrappers;

        /**
         * 自己依赖的所有
         */
        private List<DependencyWrapper> dependWrappers;

        /**
         * 存储强依赖于自己的wrapper集合
         */
        private Set<WorkerWrapper<?, ?>> selfIsMustSet;

        private boolean needCheckNextWrapperResult = true;

        public Builder<T, R> worker(IWorker<T, R> worker) {
            this.worker = worker;
            return this;
        }

        public Builder<T, R> param(T param) {
            this.param = param;
            return this;
        }

        public Builder<T, R> id(String id) {
            if (id != null) {
                this.id = id;
            }
            return this;
        }

        public Builder<T, R> needCheckNextWrapperResult(boolean needCheckNextWrapperResult) {
            this.needCheckNextWrapperResult = needCheckNextWrapperResult;
            return this;
        }

        public Builder<T, R> callback(ICallback<T, R> callback) {
            this.callback = callback;
            return this;
        }

        public Builder<T, R> depend(WorkerWrapper<?, ?>... wrappers) {
            if (wrappers == null) {
                return this;
            }
            for (WorkerWrapper<?, ?> wrapper : wrappers) {
                depend(wrapper);
            }
            return this;
        }

        public Builder<T, R> depend(WorkerWrapper<?, ?> wrapper, boolean isMust) {
            if (wrapper == null) {
                return this;
            }
            DependencyWrapper dependWrapper = new DependencyWrapper(wrapper, isMust);
            if (dependWrappers == null) {
                dependWrappers = new ArrayList<>();
            }
            dependWrappers.add(dependWrapper);
            return this;
        }

        public Builder<T, R> depend(WorkerWrapper<?, ?> wrapper) {
            return depend(wrapper, true);
        }


        public Builder<T, R> next(WorkerWrapper<?, ?> wrapper, boolean selfIsMust) {
            if (nextWrappers == null) {
                nextWrappers = new ArrayList<>();
            }
            nextWrappers.add(wrapper);

            // 强依赖自己
            if (selfIsMust) {
                if (selfIsMustSet == null) {
                    selfIsMustSet = new HashSet<>();
                }
                selfIsMustSet.add(wrapper);
            }
            return this;
        }

        public Builder<T, R> next(WorkerWrapper<?, ?> wrapper) {
            return next(wrapper, true);
        }

        public Builder<T, R> next(WorkerWrapper<?, ?>... wrappers) {
            if (wrappers == null) {
                return this;
            }
            for (WorkerWrapper<?, ?> wrapper : wrappers) {
                next(wrapper);
            }
            return this;
        }

        public WorkerWrapper<T, R> build() {
            WorkerWrapper<T, R> wrapper = new WorkerWrapper<>(id, worker, param, callback);
            wrapper.setNeedCheckNextWrapperResult(needCheckNextWrapperResult);
            if (dependWrappers != null) {
                for (DependencyWrapper workerWrapper : dependWrappers) {
                    workerWrapper.getDependWrapper().addNext(wrapper);
                    wrapper.addDepend(workerWrapper);
                }
            }
            if (nextWrappers != null) {
                for (WorkerWrapper<?, ?> workerWrapper : nextWrappers) {
                    boolean must = selfIsMustSet != null && selfIsMustSet.contains(workerWrapper);
                    workerWrapper.addDepend(wrapper, must);
                    wrapper.addNext(workerWrapper);
                }
            }

            return wrapper;
        }
    }
}
